import argparse
import torch
import logging
import sys
import os

# Add project root directory to path
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, project_root)

from watermark.watermark import WatermarkEmbedder
import json

logging.basicConfig(format='[%(asctime)s] - %(message)s',
                    datefmt='%Y/%m/%d %H:%M:%S',
                    level=logging.INFO)

def main():
    parser = argparse.ArgumentParser(description='Embed Watermark into Existing Transition Matrix')
    parser.add_argument('--matrix_path', type=str, required=True,
                       help='Path to existing transition matrix file')
    parser.add_argument('--watermark_key', type=str, required=True,
                       help='Watermark key for embedding')
    parser.add_argument('--watermark_strength', type=float, default=0.01,
                       help='Watermark embedding strength (default: 0.01)')
    parser.add_argument('--num_classes', type=int, default=None,
                       help='Number of classes (auto-detect from matrix if not provided)')
    parser.add_argument('--output_path', type=str, default=None,
                       help='Output path for watermarked matrix (default: overwrite original)')
    parser.add_argument('--backup', action='store_true',
                       help='Create backup of original matrix')
    
    args = parser.parse_args()
    
    # Check if file exists
    if not os.path.exists(args.matrix_path):
        logging.error(f"Transition matrix file does not exist: {args.matrix_path}")
        return False
    
    # Load transition matrix
    logging.info(f"Loading transition matrix: {args.matrix_path}")
    T = torch.load(args.matrix_path, weights_only=False)
    logging.info(f"Transition matrix shape: {T.shape}")
    
    # Get number of classes
    if args.num_classes is None:
        num_classes = T.shape[0]
        logging.info(f"Auto-detected number of classes: {num_classes}")
    else:
        num_classes = args.num_classes
        if num_classes != T.shape[0]:
            logging.warning(f"Warning: Provided number of classes {num_classes} does not match matrix dimension {T.shape[0]}")
            num_classes = T.shape[0]
    
    # Create backup
    if args.backup:
        backup_path = args.matrix_path + '.backup'
        torch.save(T, backup_path)
        logging.info(f"Backup created: {backup_path}")
    
    # Embed watermark
    logging.info("=" * 60)
    logging.info("Embedding watermark into transition matrix...")
    logging.info(f"Watermark key: {args.watermark_key}")
    logging.info(f"Watermark strength: {args.watermark_strength}")
    logging.info("=" * 60)
    
    embedder = WatermarkEmbedder(args.watermark_key, args.watermark_strength)
    T_watermarked = embedder.embed_watermark(T, num_classes)
    
    # Verify embedded watermark
    is_match, match_rate = embedder.verify_watermark(T_watermarked, num_classes)
    logging.info(f"Watermark verification: {'✓ Match' if is_match else '✗ No Match'} (Match rate: {match_rate:.4f})")
    
    # Save watermarked transition matrix
    if args.output_path is None:
        output_path = args.matrix_path
    else:
        output_path = args.output_path
    
    torch.save(T_watermarked, output_path)
    logging.info(f"Watermarked transition matrix saved to: {output_path}")
    
    # Save watermark information
    watermark_info = {
        'watermark_key': args.watermark_key,
        'watermark_strength': args.watermark_strength,
        'num_classes': num_classes,
        'matrix_path': output_path,
        'original_matrix_path': args.matrix_path
    }
    
    # Generate watermark info file path
    base_name = os.path.splitext(output_path)[0]
    watermark_path = base_name + '_watermark.json'
    
    with open(watermark_path, 'w') as f:
        json.dump(watermark_info, f, indent=2)
    logging.info(f"Watermark information saved to: {watermark_path}")
    
    logging.info("=" * 60)
    logging.info("Watermark embedding completed!")
    logging.info("=" * 60)
    
    return True

if __name__ == '__main__':
    main()

